package com.dc.project.util;

import com.dc.project.entity.vo.TaskResultGenerateVo;
import com.own.component.common.model.R;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import java.util.function.Supplier;

/**
 * TaskSseEmitterUtil
 *
 * @author chenxueli
 * @date 2025-03-10 13:29:00
 */
@Slf4j
@Component
public class TaskSseEmitterUtil {

    public static final Map<Long, Set<SseEmitter>> EMITTER_MAP = new ConcurrentHashMap<>();

    /**
     * 添加推送的信息
     *
     * @param taskId  任务id
     * @param emitter 推送对象
     */
    public void put(Long taskId, SseEmitter emitter) {
        var emitters = EMITTER_MAP.computeIfAbsent(taskId, k -> ConcurrentHashMap.newKeySet());
        emitters.add(emitter);
        emitter.onTimeout(emitter::complete);
        emitter.onCompletion(() -> {
            System.out.println("complete");
            // 连接关闭时的清理工作
            if (!EMITTER_MAP.containsKey(taskId)) {
                return;
            }
            var set = EMITTER_MAP.get(taskId);
            if (set.isEmpty() || set.size() == 1) {
                EMITTER_MAP.remove(taskId);
            } else {
                set.remove(emitter);
            }
        });
    }

    /**
     * 发送消息
     *
     * @param taskId    任务id
     * @param eventName 事件名称
     * @param content   内容
     */
    public void send(Long taskId, String eventName, Supplier<Object> content) {
        // 获取任务对应的所有推送对象
        var emitters = EMITTER_MAP.get(taskId);
        if (emitters == null || emitters.isEmpty()) {
            return;
        }
        var data = SseEmitter.event().name(eventName).data(content.get());
        emitters.forEach(emitter -> {
            try {
                emitter.send(data);
            } catch (Exception ex) {
                emitter.completeWithError(ex);
            }
        });
    }

    /**
     * 发送生成通知
     *
     * @param taskId   任务id
     * @param detailId 详情id
     * @param itemId   项目id
     * @param status   状态
     */
    public void sendGenerateNotice(Long taskId, Long detailId, Long itemId, int status) {
        sendGenerateNotice(taskId, detailId, itemId, status, null);
    }

    /**
     * 发送生成通知
     *
     * @param taskId   任务id
     * @param detailId 详情id
     * @param itemId   项目id
     * @param status   状态
     * @param generateList 生成的文件信息列表
     */
    public void sendGenerateNotice(
            Long taskId,
            Long detailId,
            Long itemId,
            int status,
            List<TaskResultGenerateVo> generateList
    ) {
        var map = new HashMap<String, Object>();
        map.put("taskId", taskId);
        if (detailId != null) {
            map.put("detailId", detailId);
        }
        if (itemId != null) {
            map.put("itemId", itemId);
        }
        map.put("status", status);
        if (generateList != null) {
            map.put("generateList", generateList);
        }
        send(taskId, "GENERATE_NOTICE", () -> R.success(map));
    }

}
